import numpy as np
import matplotlib.pyplot as plt
import pandas as pd


class GPR:

    def __init__(self, optimize=True):
        self.is_fit = False
        self.train_X, self.train_y = None, None
        self.params = {"l": 0.5, "sigma_f": 0.2}
        self.optimize = optimize

    def fit(self, X, y):
        self.train_X = np.asarray(X)
        self.train_y = np.asarray(y)
        self.is_fit = True

    def predict(self, X):
        if not self.is_fit:
            print("GPR Model not fit yet")
            return
        X = np.asarray(X)
        kff = self.kernel(self.train_X, self.train_X)
        kyy = self.kernel(X, X)
        kfy = self.kernel(self.train_X, X)
        kff_inv = np.linalg.inv(kff + 1e-8 * np.eye(len(self.train_X)))

        mu = kfy.T.dot(kff_inv).dot(self.train_y)
        conv = kyy - kfy.T.dot(kff_inv).dot(kfy)
        return mu, conv

    def kernel(self, x1, x2):
        dist_matrix = np.sum(x1 ** 2, 1).reshape(-1, 1) + np.sum(x2 ** 2, 1) - 2 * np.dot(x1, x2.T)
        return self.params["sigma_f"] ** 2 * np.exp(-0.5 / self.params["l"] ** 2 * dist_matrix)


def y(x, noise_sigma=0.0):
    x = np.asarray(x)
    y = np.cos(x) + np.random.normal(0, noise_sigma, size=x.shape)
    return y.tolist()

# 导入 bitcoin数据 BCHAIN-MKPRU.csv
bcc_df = pd.read_csv('BCHAIN-MKPRU.csv')

# 打印预览
print(bcc_df)

# 导入 gold数据 LBMA-GOLD.csv
gold_df = pd.read_csv('LBMA-GOLD.csv')

# 打印预览
print(gold_df)
# 合并两项数据
data = pd.merge(bcc_df, gold_df, on='Date', how='outer')
data["USD (PM)"].fillna(method='ffill', inplace=True)
data["USD (PM)"].fillna(method='ffill', inplace=True)
data["USD (PM)"].fillna(method='bfill', inplace=True)

data.reset_index()

train_X = np.array(data.index).reshape(-1, 1)

train_y = np.array(data['Value']).reshape(-1, 1)
# train_X = np.array([3, 1, 4, 5, 9]).reshape(-1, 1)
# train_y = y(train_X, noise_sigma=1e-4)
test_X = np.arange(10000, 100000, 1).reshape(-1, 1)

gpr = GPR()
gpr.fit(train_X, train_y)
mu, cov = gpr.predict(test_X)
test_y = mu.ravel()
uncertainty = 1.96 * np.sqrt(np.diag(cov))  # 1.96表示95%的置信度
plt.figure()
plt.title("l=%.2f sigma_f=%.2f" % (gpr.params["l"], gpr.params["sigma_f"]))
plt.fill_between(test_X.ravel(), test_y + uncertainty, test_y - uncertainty, alpha=0.1)
plt.plot(test_X, test_y, label="predict")
plt.scatter(train_X, train_y, label="train", c="red", marker="x")
plt.legend()
plt.show()